#!/usr/bin/python3
# -*- coding: utf-8 -*-

import cv2
import numpy as np

np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

import rospy
import message_filters

from geometry_msgs.msg import Transform,Pose,Point32

from sensor_msgs.msg import Image
from cv_bridge import CvBridge, CvBridgeError


from bs_imu import bs_imu
from bs_imu import bs_gps_raw
from bs_img import bs_img_real
from bs_img import bs_img_base
from bs_img.bs_cfg_real_5m import *
from bs_img.bs_lie_solve import *
from bs_img.serial_pack import *

import time
'''
输入图像
输出相对的解算值
'''




class BS_Solver:
    def __init__(self):
        sf_img_topic = "/camera_sf/image_dh"
        self.last_plane_pts = np.array([])   #保存上一帧点用于配对
        self.is_first_frame = True                  #判断是否为第一帧

        self.img_bridge = CvBridge()

        self.invalid_frame_num_plane = 0      #compute continus invalid frame
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 8     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_plane = np.array([np.nan, np.nan, np.nan])
        self.w_plane = 1.5
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.w_drogue = 1.5

        self.sf_img_sub = rospy.Subscriber(sf_img_topic, Image, self.sf_img_cb)

        self.sf_plane_pub = rospy.Publisher("/bs_debuger/result/sf_plane", Transform , queue_size=1)
        self.sf_drogue_pub = rospy.Publisher("/bs_debuger/result/sf_drogue", Transform, queue_size=1)

        # self.ts = message_filters.TimeSynchronizer([self.lf_img_sub, self.sf_img_sub], queue_size=10)
        # self.ts.registerCallback(self.img_cb)



    def sf_img_cb(self,img_msg):
        img_gray = self.img_bridge.imgmsg_to_cv2(img_msg, img_msg.encoding)
        # img_gray = cv2.undistort(img_gray,camK_sf, distortion_coeffs_sf)  #time-consuming/changed to solvePnP
        _, img_bin = cv2.threshold(img_gray, 145, 0xff, cv2.THRESH_BINARY)
        # 如果相机倒置需要旋转一下
        # lf_img_gray = cv2.rotate(lf_img_gray, cv2.ROTATE_180)
        plane_t_vec = [np.nan]*3
        plane_quat = [np.nan]*4
        drogue_t_vec = [np.nan]*3
        
        gt_pts, gt_pts_plane = bs_img_real.img_to_pts_new(img_bin)
        cnt_plane_pts = []

        print(gt_pts_plane)
        
        if len(gt_pts) >= 11:
            if len(gt_pts_plane)>=5:
                plane_result_dict = bs_img_base.solve_plane_pt(gt_pts_plane,plane_real_ptL,cam_K=camK_sf,distort_coeffs=distortion_coeffs_sf)
                if plane_result_dict["img_valid"]:
                    self.is_first_frame = False 
                    cnt_plane_pts = plane_result_dict["plane_px_calcL"]
                    plane_t_vec = plane_result_dict["t_vec"].flatten()
                    plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())
            # 获取剩下的点用来拟合锥套
            set_plane = set(map(tuple, cnt_plane_pts))
            set_all = set(map(tuple, np.array(gt_pts)))
            set_drogue = set_all - set_plane
            gt_pts_drogue = [np.array(item) for item in set_drogue] 
            drogue_result_dict = bs_img_real.solve_drogue(gt_pts_drogue, camK_sf, taper_radius=r_drogue)
                
            drogue_t_vec = drogue_result_dict["t_vec"]

        else:
            if len(gt_pts)>=5:
                plane_result_dict = bs_img_base.solve_plane_pt(gt_pts,plane_real_ptL,cam_K=camK_sf,distort_coeffs=distortion_coeffs_sf)
                if plane_result_dict["img_valid"]:
                    self.is_first_frame = False 
                    cnt_plane_pts = plane_result_dict["plane_px_calcL"]
                    plane_t_vec = plane_result_dict["t_vec"].flatten()
                    plane_quat = bs_imu.rot_vec_to_quat(plane_result_dict["r_vec"].flatten())

            # 获取剩下的点用来拟合锥套
            set_plane = set(map(tuple, cnt_plane_pts))
            set_all = set(map(tuple, np.array(gt_pts)))
            set_drogue = set_all - set_plane
            drogue_pts = [np.array(item) for item in set_drogue]  
            drogue_result_dict = bs_img_real.solve_drogue(drogue_pts, camK_sf, taper_radius=r_drogue)
                
            drogue_t_vec = drogue_result_dict["t_vec"]
        
            
        
        # KF for plane
        plane_is_visiable = True if not(np.all(np.isnan(plane_t_vec))) else False
        if not(plane_is_visiable):
            self.invalid_frame_num_plane = self.invalid_frame_num_plane + 1
            if self.invalid_frame_num_plane <= self.invalid_frame_max:
                self.plane_res = self.x_plane
                if not(np.all(np.isnan(self.plane_res))):
                    plane_res_frame = serial_frame("plane", True, round((self.plane_res[2]-0.6)*100),round(self.plane_res[0]*100),round((self.plane_res[1]-0.13)*100))  
                else:
                    plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
            else:
                plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
        
        else:
            self.invalid_frame_num_plane = 0

            if np.all(np.isnan(self.x_plane)):   #the first frame
                self.x_plane = plane_t_vec
            else:
                miu = plane_t_vec-self.x_plane
                self.x_plane = self.x_plane + np.matmul(self.Q,np.matmul(np.linalg.pinv(self.Q+self.R/self.w_plane),miu.T))
                self.w_plane = 1.5/(1+np.matmul(miu,np.matmul(np.linalg.pinv(self.R),miu.T)))
            plane_res_frame = serial_frame("plane",True,round((self.x_plane[2])*100),round(self.x_plane[0]*100),round((self.x_plane[1])*100))  

        # KF for drogue
        drogue_is_visiable = True if not(np.all(np.isnan(drogue_t_vec))) else False
        if not(drogue_is_visiable):
            self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
            if self.invalid_frame_num_drogue <= self.invalid_frame_max:
                self.drogue_res = self.x_drogue
                if not(np.all(np.isnan(self.drogue_res))):
                    drogue_res_frame = serial_frame("drogue", True, round((self.drogue_res[2])*100),round(self.drogue_res[0]*100),round((self.drogue_res[1])*100))  
                else:
                    drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
            else:
                drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
        
        else:
            self.invalid_frame_num_drogue = 0

            if np.all(np.isnan(self.x_drogue)):   #the first frame
                self.x_drogue = drogue_t_vec
            else:
                miu = drogue_t_vec-self.x_drogue
                self.x_drogue = self.x_drogue + np.matmul(self.Q,np.matmul(np.linalg.pinv(self.Q+self.R/self.w_drogue),miu.T))
                self.w_drogue = 1.5/(1+np.matmul(miu,np.matmul(np.linalg.pinv(self.R),miu.T)))
            drogue_res_frame = serial_frame("drogue",True,round((self.x_drogue[2])*100),round(self.x_drogue[0]*100),round((self.x_drogue[1])*100))  

        plane_msg = Transform()
        plane_msg.translation.x = plane_t_vec[0]
        plane_msg.translation.y = plane_t_vec[1]
        plane_msg.translation.z = plane_t_vec[2]
        # plane_msg.rotation.w = plane_quat[0]
        # plane_msg.rotation.x = plane_quat[1]
        # plane_msg.rotation.y = plane_quat[2]
        # plane_msg.rotation.z = plane_quat[3]
        self.sf_plane_pub.publish(plane_msg)
        
        drogue_msg = Transform()
        drogue_msg.translation.x = drogue_t_vec[0]
        drogue_msg.translation.y = drogue_t_vec[1]
        drogue_msg.translation.z = drogue_t_vec[2]
        self.sf_drogue_pub.publish(drogue_msg)




if __name__ == "__main__":
    rospy.init_node('bs_solver_node', anonymous=True)
    bs_solver = BS_Solver()
    rospy.spin()
